from torch.utils.data import DataLoader
from json import load
from os import get_terminal_size
from rich import print
from .name_class import (
    NameClassDataset,
    LSTMRebuild,
    Model
)

DataType = tuple[list[str], list[str]]

def separate_line(sep: str, style: str) -> None:
    """Separate line with style with length of terminal window"""
    terminal_width = get_terminal_size().columns
    print(f'[{style}]' + sep * terminal_width + '[/]')

def read_data(file_name: str, encoding: str='utf-8') -> DataType:
    """Read data from txt file and return a tuple of names and countries"""
    names, countries = [], []
    with open(file_name, 'r', encoding=encoding) as f:
        for line in f.readlines():
            if len(line) <= 5:
                continue
            (name, country) = line.strip().split('\t')
            names.append(name)
            countries.append(country)
    return names, countries

def run_name_class(file_name: str, model: Model,
                   head: int=5, **kwargs) -> None:
    """
    Run a name classification model

    Args:
        file_name: The path to the name classification data
        model: The model to run, either `RNN`, `LSTM` or `GRU`
        head: Display the first `head` results
    """
    names, countries = read_data(file_name)
    name_class = NameClassDataset(names, countries)

    model_rebuild = model(**kwargs)
    dl = DataLoader(name_class, batch_size=1, shuffle=True)

    print('[bold magenta]RNN model:[/]')
    for i, (name, _) in enumerate(dl):
        if type(model_rebuild) == LSTMRebuild:
            hidden, c = model_rebuild.init_hidden()
            outputs, hidden, c = model_rebuild(name[0], hidden, c)
        else:
            hidden = model_rebuild.init_hidden()
            outputs, hidden = model_rebuild(name[0], hidden)

        print(f'[bold green]output:[/]\n{outputs}')
        print(f'[bold green]output shape:[/] {list(outputs.shape)}')
        print(f'[bold green]hidden:[/]\n{hidden}')
        print(f'[bold green]hidden shape:[/] {list(hidden.shape)}')
        if i == head - 1:
            break

def read_json(json_path: str, encoding: str='utf-8') -> dict:
    """Read json file and reeturn dictionary"""
    with open(json_path, 'r', encoding=encoding) as json_f:
        return load(json_f)

def human_readable(size: int) -> str:
    """Convert bytes size to human readable format"""
    unit = ['B', 'KB', 'MB', 'GB', 'TB']
    idx = 0
    while size > 1024 and idx < len(unit) - 1:
        size /= 1024
        idx += 1
    return '%.1f%s' % (size, unit[idx])
